/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at

* http://www.apache.org/licenses/LICENSE-2.0

* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.

* File sample_process.cpp
* Description: handle acl resource
*/
#include "detectPostprocess.h"

#include <iostream>

#include "AclLiteApp.h"
#include "AclLiteUtils.h"
#include "CarParams.h"
#include "acl/acl.h"

using namespace std;

namespace {
const string kCarColorClass[9] = {"black", "blue",   "brown", "green", "pink",
                                  "red",   "silver", "white", "yellow"};
const uint32_t kLineSolid = 2;
const uint32_t kLabelOffset = 11;
const double kFountScale = 0.5;
const cv::Scalar kFontColor(0, 0, 255);
const vector<cv::Scalar> kColors{
    cv::Scalar(237, 149, 100), cv::Scalar(0, 215, 255), cv::Scalar(50, 205, 50),
    cv::Scalar(139, 85, 26)};

const uint32_t kConfidenceBufId = 0;
const uint32_t kBboxBufId = 1;
uint32_t kModelWidth = 960;
uint32_t kModelHeight = 960;
float kConfidence = 0.2;
float kNMS = 0.6;
}  // namespace

struct DBBox {
  cv::Rect rect;
  float score;
};

float doIou(const cv::Rect &a, const cv::Rect &b) {
  cv::Rect inter = a | b;
  cv::Rect combine = a & b;
  return static_cast<float>(combine.area()) / inter.area();
}

std::vector<DBBox> doNms(std::vector<DBBox> &bbox, float nms) {
  std::sort(bbox.begin(), bbox.end(),
            [](const DBBox &a, const DBBox &b) { return a.score > b.score; });

  std::vector<DBBox> result;
  std::vector<bool> available(bbox.size(), true);
  for (size_t i = 0; i < bbox.size(); i++) {
    if (available[i] == false) continue;

    result.push_back(bbox[i]);
    for (size_t j = i + 1; j < bbox.size(); j++) {
      if (available[j] == false) continue;
      float iou = doIou(bbox[i].rect, bbox[j].rect);
      if (iou > nms) {
        available[j] = false;
      }
    }
  }
  return result;
}

std::vector<DBBox> DecodeBBox(float *confidence_feature, float *box_feature,
                              int32_t feature_size,
                              const std::vector<float> &scales,
                              float confidence, float nms) {
  std::vector<DBBox> filter_bboxes;
  for (uint32_t i = 0; i < feature_size; i++) {
    if (confidence_feature[i] > confidence) {
      int xmin = box_feature[i * 4 + 0] / scales[0];
      int ymin = box_feature[i * 4 + 1] / scales[1];
      int xmax = box_feature[i * 4 + 2] / scales[0];
      int ymax = box_feature[i * 4 + 3] / scales[1];
      DBBox bbox{cv::Rect(cv::Point(xmin, ymin), cv::Point(xmax, ymax)),
                 confidence_feature[i]};
      filter_bboxes.emplace_back(bbox);
    }
  }

  return doNms(filter_bboxes, nms);
}

DetectPostprocessThread::DetectPostprocessThread() {}

DetectPostprocessThread::~DetectPostprocessThread() {}

AclLiteError DetectPostprocessThread::Init() { return ACLLITE_OK; }

AclLiteError DetectPostprocessThread::Process(int msgId,
                                              shared_ptr<void> data) {
  AclLiteError ret = ACLLITE_OK;
  switch (msgId) {
    case MSG_DETECT_INFER_OUTPUT:
      InferOutputProcess(static_pointer_cast<CarDetectDataMsg>(data));
      MsgSend(static_pointer_cast<CarDetectDataMsg>(data));
      break;
    case MSG_ENCODE_FINISH:
      SendMessage(kMainThreadId, MSG_APP_EXIT, nullptr);
      break;
    default:
      ACLLITE_LOG_INFO("Detect PostprocessThread thread ignore msg %d", msgId);
      break;
  }

  return ret;
}

AclLiteError DetectPostprocessThread::MsgSend(
    shared_ptr<CarDetectDataMsg> carDetectDataMsg) {
  while (1) {
    AclLiteError ret = SendMessage(carDetectDataMsg->detectPostThreadId,
                                   MSG_ENCODE_FINISH, carDetectDataMsg);
    if (ret == ACLLITE_ERROR_ENQUEUE) {
      usleep(500);
      continue;
    } else if (ret == ACLLITE_OK) {
      break;
    } else {
      ACLLITE_LOG_ERROR("Send read frame message failed, error %d", ret);
      return ret;
    }
  }

  return ACLLITE_OK;
}

AclLiteError DetectPostprocessThread::DrawResultOnPic(
    shared_ptr<CarDetectDataMsg> &carDetectDataMsg) {
  for (int i = 0; i < carDetectDataMsg->carInfo.size(); ++i) {
    cv::rectangle(carDetectDataMsg->frame,
                  carDetectDataMsg->carInfo[i].rectangle.lt,
                  carDetectDataMsg->carInfo[i].rectangle.rb,
                  kColors[i % kColors.size()], kLineSolid);
    cv::putText(
        carDetectDataMsg->frame, carDetectDataMsg->carInfo[i].detect_result,
        cv::Point(carDetectDataMsg->carInfo[i].rectangle.lt.x - kLabelOffset,
                  carDetectDataMsg->carInfo[i].rectangle.lt.y - kLabelOffset),
        cv::FONT_HERSHEY_COMPLEX, kFountScale, kFontColor);
    cv::putText(
        carDetectDataMsg->frame, carDetectDataMsg->carInfo[i].carColor_result,
        cv::Point(carDetectDataMsg->carInfo[i].rectangle.lt.x,
                  carDetectDataMsg->carInfo[i].rectangle.lt.y + kLabelOffset),
        cv::FONT_HERSHEY_COMPLEX, kFountScale, kFontColor);
  }

  stringstream sstream;
  sstream.str("");
  sstream << "../out/output/device_" << carDetectDataMsg->channelId
          << "_out_pic_" << carDetectDataMsg->frameNum << ".jpg";
  cv::imwrite(sstream.str(), carDetectDataMsg->frame);
  return ACLLITE_OK;
}

AclLiteError DetectPostprocessThread::InferOutputProcess(
    shared_ptr<CarDetectDataMsg> carDetectDataMsg) {
  if (carDetectDataMsg->isLastFrame) return ACLLITE_OK;

  float *detectConfidencePtr =
      (float *)carDetectDataMsg->detectInferData[kConfidenceBufId].data.get();
  float *detectBboxPtr =
      (float *)carDetectDataMsg->detectInferData[kBboxBufId].data.get();

  float scale_w = carDetectDataMsg->resizedFrame.width * 1.0f /
                  carDetectDataMsg->frame.cols;
  float scale_h = carDetectDataMsg->resizedFrame.height * 1.0f /
                  carDetectDataMsg->frame.rows;

  auto detectResult = DecodeBBox(
      detectConfidencePtr, detectBboxPtr,
      carDetectDataMsg->detectInferData[kConfidenceBufId].size / sizeof(float),
      {scale_w, scale_h}, kConfidence, kNMS);

  ACLLITE_LOG_INFO("Total result: %ld", detectResult.size());
  for (auto &bbox : detectResult) {
    CarInfo carInfo;
    carInfo.rectangle.lt = bbox.rect.tl();
    carInfo.rectangle.rb = bbox.rect.br();
    carInfo.detect_result =
        "elec_cycle" + std::to_string(bbox.score * 100) + "\%";
    ACLLITE_LOG_INFO(" score: %f", bbox.score);
    carDetectDataMsg->carInfo.emplace_back(carInfo);
  }

  auto ret = DrawResultOnPic(carDetectDataMsg);
  if (ret != ACLLITE_OK) {
    ACLLITE_LOG_ERROR("Send image to presentAgent failed, error %d", ret);
    return ACLLITE_ERROR;
  }
  return ACLLITE_OK;
}
